use nalgebra::vector;
use nalgebra::{Vector2, Matrix2};


pub struct Fem {
    nw: usize,
    volumn0: f32,
    pub pxs: Vec<Vector2<f32>>,
    pvs: Vec<Vector2<f32>>,
    fbs: Vec<Matrix2<f32>>,
    pfs: Vec<Vector2<f32>>,
    container: Object,

    pub cursor_world_pos: Vector2<f32>,
    pub cursor_force: f32,
}

impl Fem {
    pub fn new() -> Self {

        let w = 0.3;
        let h = 0.1;
        let dx = 0.02;
        // number of segment
        let nw = (w / dx) as usize;
        let nh = (h / dx) as usize;
        let nf = nw * nh * 2;
        // let nv = (nw + 1) * (nh + 1);

        // let ball_pos: Vector2<f32> = vector![0.5f32, 0.0];
        // let ball_r = 0.32;
        // let damping = 12.5;

        let mut pxs: Vec<Vector2<f32>> = Vec::new();
        let min_pos: Vector2<f32> = vector![0.45f32, 0.45];
        for j in 0..nh + 1 {
            for i in 0..nw + 1 {
                let pos: Vector2<f32> = min_pos + vector!(i as f32, j as f32) * dx;
                pxs.push(pos);
            }
        }
        let pvs: Vec<Vector2<f32>> = vec![vector!(0.0f32, 0.0); pxs.len()];
        let pfs: Vec<Vector2<f32>> = vec![Vector2::zeros(); pxs.len()];

        let mut fbs: Vec<Matrix2<f32>> = Vec::with_capacity(nf);
        for i in 0..nf {
            let [i1, i2, i3] = f2v(i, nw);
            let [x1, x2, x3] = [pxs[i1], pxs[i2], pxs[i3]];
            let b_inv = Matrix2::from_columns(&[x1 - x3, x2 - x3]);
            fbs.push(b_inv.try_inverse().unwrap())
        }
        let volumn0 = {
            let [i1, i2, i3] = f2v(0, nw);
            let [x1, x2, x3] = [pxs[i1], pxs[i2], pxs[i3]];
            let b_inv = Matrix2::from_columns(&[x1 - x3, x2 - x3]);
            b_inv.determinant() / 6.0
        };
        assert!(volumn0 > 0.0);

        let container = Object::new(
            vector![0.5, 0.5],
            vector![1.0, 1.0]
        );

        Self {
            nw, pxs, pvs, fbs, pfs, volumn0, container,
            cursor_force: 0.0, cursor_world_pos: vector![0.0, 0.0],
        }
    }
    pub fn step(&mut self, dt: f32) {

        // fem forces
        let young = 4e6;
        let poisson = 0.2;
        let mu = young * 0.5 / (1.0 + poisson);
        let lam = young * poisson / (1.0 + poisson) / (1.0 - 2.0 * poisson);
        self.pfs.fill(Vector2::zeros());
        for i in 0..self.fbs.len() {
            let [i1, i2, i3] = f2v(i, self.nw);
            let [x1, x2, x3] = [self.pxs[i1], self.pxs[i2], self.pxs[i3]];
            let d = Matrix2::from_columns(&[x1 - x3, x2 - x3]);
            let f = d * self.fbs[i];
            let p = neohookean(f, mu, lam);
            let h = -self.volumn0 * p * self.fbs[i].transpose();
            let [h1, h2] = [h.column(0), h.column(1)];
            self.pfs[i1] += h1;
            self.pfs[i2] += h2;
            self.pfs[i3] -= h1 + h2;
        }

        // motion and border boundary
        let bounce = 1.0;
        let minx = self.container.pos - 0.5 * self.container.size;
        let maxx = self.container.pos + 0.5 * self.container.size;
        for i in 0..self.pxs.len() {
            // elastic force
            self.pvs[i] += self.pfs[i] * dt;
            // gravity
            self.pvs[i].y -= 9.8 * dt;
            // cursor force
            let cursor_offset = self.cursor_world_pos - self.pxs[i];
            let cursor_distance = cursor_offset.magnitude();
            if cursor_distance < 0.2 {
                self.pvs[i] += cursor_offset / cursor_distance * self.cursor_force * dt;
            }
            self.pvs[i] = self.pvs[i] - self.pvs[i] * self.pvs[i].magnitude() * 0.3 * dt;

            self.pxs[i] += self.pvs[i] * dt;

            if self.pxs[i].y < minx.y {
                self.pxs[i].y = minx.y;
                self.pvs[i].y = -self.pvs[i].y * bounce;
            } else if self.pxs[i].y > maxx.y {
                self.pxs[i].y = maxx.y;
                self.pvs[i].y = -self.pvs[i].y * bounce;
            }
            if self.pxs[i].x < minx.x {
                self.pxs[i].x = minx.x;
                self.pvs[i].x = -self.pvs[i].x * bounce;
            } else if self.pxs[i].x > maxx.x {
                self.pxs[i].x = maxx.x;
                self.pvs[i].x = -self.pvs[i].x * bounce;
            }
        }
    }
}

fn f2v(f: usize, nw: usize) -> [usize; 3] {
    let fh = f / 2;
    let fl = f % 2;
    let j = fh / nw;
    let i = fh % nw;

    let i1 = i + 1;
    let j1 = j + 1;
    let nw1 = nw + 1;
    let b = i1 + j * nw1;
    let c = i + j1 * nw1;
    let ad = i + fl + (j + fl) * nw1;
    return [ad, b, c]
}

fn neohookean(f: Matrix2<f32>, mu: f32, lambda: f32) -> Matrix2<f32> {
    let f_it = f.try_inverse().unwrap().transpose();
    let j = f.determinant();
    return mu * (f - f_it) + lambda * j.ln() * f_it;
}

struct Object {
    pos: Vector2<f32>,
    size: Vector2<f32>,
}
impl Object {
    fn new(pos: Vector2<f32>, size: Vector2<f32>) -> Self {
        Self { pos, size }
    }
}